import logging
import numpy as np

from util.util_log import setup_logging
from FitnessMultiBeamTrace import FitnessMultiBeamTrace


# 配置日志，默认打印到控制台，也可以设置打印到文件
setup_logging()
# setup_logging(log_file="../files/multi-beam-trace/beam2/QE/L=L1+L4/trace_2025-03-19-1644-16beam-qe.txt")
# 获取日志记录器并记录日志
logger = logging.getLogger("[BGA]")


class BinaryGeneticAlgorithm:
    def __init__(self, objective_function, population_size, dimensions, mutation_rate=0.01, generations=100):
        """
        初始化二进制遗传算法。

        :param objective_function: 适应度函数
        :param population_size: 种群大小
        :param dimensions: 搜索空间维度 (tuple of two integers for 2D arrays)
        :param mutation_rate: 变异率
        :param generations: 迭代次数
        """
        self.objective_function = objective_function
        self.population_size = population_size
        self.dimensions = dimensions
        self.mutation_rate = mutation_rate
        self.generations = generations
        self.population = []
        self.best_individual = None
        self.best_fitness = float('-inf')

    def __generate_population_part(self, before_array, current_array, max_num):
        part = []
        # 添加before和current作为前两个元素
        part.append(before_array.copy())
        part.append(current_array.copy())
        # 计算 diff 数组
        diff = before_array != current_array
        # 获取所有 False 的坐标
        false_coords = np.argwhere(diff)
        # 生成剩余的个体
        for _ in range(max_num - 2):
            # 复制 current_array
            modified_array = np.copy(current_array)
            # 随机选择一些差异行
            if np.sum(diff == True) > 0:
                # 随机选择一部分 False 坐标
                num_to_select = np.random.randint(0, len(false_coords) + 1)
                selected_false_coords = false_coords[np.random.choice(len(false_coords), num_to_select, replace=False)]
                # 将 modified_array 中选取坐标替换为 before_array
                for coord in selected_false_coords:
                    row, col = coord
                    modified_array[row, col] = before_array[row, col]
                # print("Modified array:", modified_array.tolist())
            part.append(modified_array)
        return part

    def initialize_population(self):
        """初始化种群"""
        self.population = [np.random.randint(2, size=self.dimensions) for _ in range(self.population_size)]

    def initialize_population_group(self, phaseBit_mix_ps_before, phaseBit_mix_ps, phaseBit_mix_qe_before,
                                    phaseBit_mix_qe):
        parts_ps = self.__generate_population_part(phaseBit_mix_ps_before, phaseBit_mix_ps, self.population_size // 2)
        parts_qe = self.__generate_population_part(phaseBit_mix_qe_before, phaseBit_mix_qe,
                                                 self.population_size - self.population_size // 2)
        self.population = parts_ps + parts_qe

    def selection(self):
        """选择操作：轮盘赌选择"""
        fitness_values = np.array([self.objective_function(individual) for individual in self.population])
        probabilities = fitness_values / fitness_values.sum()
        selected_indices = np.random.choice(range(self.population_size), size=self.population_size, p=probabilities)
        selected_population = [self.population[i] for i in selected_indices]
        return selected_population

    def crossover(self, parent1, parent2):
        """交叉操作：单点交叉"""
        if np.random.rand() < 0.5:
            crossover_point = np.random.randint(1, np.prod(self.dimensions) - 1)
            child1 = np.concatenate((parent1[:crossover_point], parent2[crossover_point:]))
            child2 = np.concatenate((parent2[:crossover_point], parent1[crossover_point:]))
        else:
            child1 = np.copy(parent1)
            child2 = np.copy(parent2)
        return child1.reshape(self.dimensions), child2.reshape(self.dimensions)

    def mutate(self, individual):
        """变异操作"""
        flat_individual = individual.flatten()
        for i in range(len(flat_individual)):
            if np.random.rand() < self.mutation_rate:
                flat_individual[i] = 1 - flat_individual[i]
        return flat_individual.reshape(self.dimensions)

    def run(self, phaseBit_mix_ps_before, phaseBit_mix_ps, phaseBit_mix_qe_before, phaseBit_mix_qe):
        """运行二进制遗传算法"""

        # """初始化返回值"""
        self.best_individual = phaseBit_mix_ps
        self.best_fitness = self.objective_function(self.best_individual)
        self.best_individual_history = []
        self.best_fitness_history = []

        # # 普通初始化
        # self.initialize_population()
        # 初始化
        self.initialize_population_group(phaseBit_mix_ps_before, phaseBit_mix_ps,
                                         phaseBit_mix_qe_before, phaseBit_mix_qe)

        for generation in range(self.generations):
            new_population = []

            # 选择操作
            selected_population = self.selection()

            # 交叉和变异操作
            for i in range(0, self.population_size, 2):
                parent1 = selected_population[i].flatten()
                parent2 = selected_population[i + 1].flatten()
                child1, child2 = self.crossover(parent1, parent2)
                child1 = self.mutate(child1)
                child2 = self.mutate(child2)
                new_population.append(child1)
                new_population.append(child2)

            # 更新种群
            self.population = new_population

            # 找到当前最佳个体
            for individual in self.population:
                fitness_value = self.objective_function(individual)
                if fitness_value > self.best_fitness:
                    self.best_fitness = fitness_value
                    self.best_individual = individual

            # 记录最佳适应度曲线
            self.best_fitness_history.append(self.best_fitness)
            self.best_individual_history.append(self.best_individual)

            # 打印当前最佳适应度
            # print(f"Generation {generation + 1}: Best Fitness = {self.best_fitness}")
            logger.info("generation=%d: self.best_fitness=%f" % (generation, self.best_fitness))
            logger.info("---------------------------------------------------------------------------")

        # return self.best_individual, self.best_fitness
        return self.best_individual, self.best_fitness, self.best_fitness_history, self.best_individual_history



def test():
    # 使用二进制遗传算法优化最大化一的数量的问题
    dimensions = (3, 3)  # Example 3x3 matrix
    population_size = 30
    mutation_rate = 0.01
    generations = 100

    # 创建一个示例 FitnessMultiBeamTrace 实例
    bit_num = 8  # 根据实际情况设置比特数
    beam_num = 2  # 根据实际情况设置波束数
    fitness_function = FitnessMultiBeamTrace(bit_num, beam_num)

    # 初始化主瓣位置
    phaseBit_target = np.random.randint(2, size=(3, 3))
    fitness_function.init_main_lope_pos(phaseBit_target)

    # 定义初始相位比特矩阵
    phaseBit_ps_before = np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]])
    phaseBit_ps = np.array([[0, 0, 0], [0, 0, 0], [1, 1, 1]])
    phaseBit_qe_before = np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]])
    phaseBit_qe = np.array([[1, 1, 1], [0, 0, 0], [0, 0, 0]])

    # 使用自定义的适应度函数
    bga_custom_fitness = BinaryGeneticAlgorithm(objective_function=fitness_function.fitness,
                                                population_size=population_size,
                                                dimensions=dimensions,
                                                mutation_rate=mutation_rate,
                                                generations=generations)
    best_individual_custom, best_fitness_custom, best_fitness_history_custom, best_individual_history_custom \
        = bga_custom_fitness.run(phaseBit_ps_before, phaseBit_ps, phaseBit_qe_before, phaseBit_qe)

    print(f"\nFinal Best Individual with Custom Fitness Function:\n{best_individual_custom}")
    print(f"Final Best Fitness with Custom Fitness Function: {best_fitness_custom}")


if __name__=="__main__":
    test()